import os
import glob
import pdb
import argparse
import numpy as np
from pathlib import Path
from scipy.io import loadmat,savemat

COMBINE_COEFF = False
COMBINE_COEFF_RBGFACE = True

COMBINE_LAMDK = False
COMBINE_LAMDK_RBGFACE = True

COMBINE_BBOX  = False


def arg_parser():
    parser = argparse.ArgumentParser("Extract video frames")
    parser.add_argument("--path_frames", type=str, help="video path for extract")
    parser.add_argument("--path_data", type=str, help='save extraction frames')
    parser.add_argument("--override", action='store_true', help="if override origin frames")

    args = parser.parse_args()
    return args


if __name__ == '__main__':
    
    args = arg_parser()

    path_data = Path(args.path_data)
    path_frames = path_data.joinpath("frames")
    
    v_frames = sorted(path_frames.glob("*.jpg"), key=lambda x: int(x.stem[-4:]))
    print(f"contain {len(v_frames)} frames in {str(path_data)}")
    
    # combine 3dmm
    if COMBINE_COEFF:

        jpath_coeffs = path_data.joinpath("reconstruction/coeff_mat")
        v_coeffs = sorted(jpath_coeffs.glob("*.mat"))
        print(f"contain {len(v_coeffs)} reconstruction in {str(path_data)}")

        if len(v_coeffs) != len(v_frames):
            print(f"\t ERROR && Continue: This video missing coeff, {str(jpath_coeffs)} ")
            # continue
        else:
            coeff_all = np.zeros((len(v_frames), 257), np.float32)
            for k in range(len(v_frames)):
                # load mat data
                kpath_mat = jpath_coeffs.joinpath(f"frame-{k:04d}.mat")
                coeff_all[k, :] = loadmat(str(kpath_mat))['coeff']

            # save combine data
            jpath_coeff_save = path_data.joinpath("video/coeff_deepconstruction.npy")
            jpath_coeff_save.parent.mkdir(parents=True, exist_ok=True)
            np.save(str(jpath_coeff_save), coeff_all)
            print(f"combine the coeff: {str(jpath_coeffs)} --> {str(jpath_coeff_save)}")

    # combine landmark
    if COMBINE_LAMDK:
        
        jpath_landmk = path_data.joinpath("landmk")
        v_landmk = sorted(jpath_landmk.glob("*.txt"), key=lambda x: int(x.stem[-4:]))
        print(f"contain {len(v_landmk)} landmark in {str(path_data)}")

        if len(v_landmk) != len(v_frames):
            print(f"\t ERROR && Continue: This video missing landmrk, {str(jpath_landmk)} ")
            # continue
        else:
            landmk_all = np.zeros((len(v_frames), 68, 2), np.int)
            for k in range(len(v_frames)):
                # load mat data
                kpath_txt = jpath_landmk.joinpath(f"frame-{k:04d}.txt")
                landmk_all[k, :] = np.loadtxt(str(kpath_txt), dtype=int, delimiter=" ")

            # print(f"landmk_all shape = {landmk_all.shape}, = {landmk_all[0]}")
            # save combine data
            jpath_coeff_save = path_data.joinpath("video/landmark_dlib.npy")
            jpath_coeff_save.parent.mkdir(parents=True, exist_ok=True)
            np.save(str(jpath_coeff_save), landmk_all)
            print(f"combine the coeff: {str(jpath_landmk)} --> {str(jpath_coeff_save)}")

    # combine landmark bbox
    if COMBINE_BBOX:
        jpath_bbox = path_data.joinpath("landmk_bbox")
        v_bbox = sorted(jpath_bbox.glob("*.txt"), key=lambda x: int(x.stem[-4:]))
        print(f"contain {len(v_bbox)} bbox in {str(path_data)}")

        if len(v_bbox) != len(v_frames):
            print(f"\t ERROR && Continue: This video missing bbox, {str(jpath_bbox)} ")
            # continue
        else:
            bbox_all = np.zeros((len(v_frames), 4), np.int)
            
            for k in range(len(v_frames)):
                # load mat data
                kpath_txt = jpath_bbox.joinpath(f"frame-{k:04d}.txt")
                bbox_all[k] = np.loadtxt(str(kpath_txt), dtype=int, delimiter=" ").reshape(-1)

            # print(f"bbox_all shape = {bbox_all.shape}, = {bbox_all[0]}")
            # save combine data
            jpath_bbox_save = path_data.joinpath("video/landmark_bbox.npy")
            jpath_bbox_save.parent.mkdir(parents=True, exist_ok=True)
            np.save(str(jpath_bbox_save), bbox_all)
            print(f"combine the bbox: {str(jpath_bbox)} --> {str(jpath_bbox_save)}")


    # combine landmark of RGBFaceNet
    if COMBINE_LAMDK_RBGFACE:
        jpath_ladmk_proj = path_data.joinpath("rgbface/landmark_proj")
        jpath_ladmk_crop = path_data.joinpath("rgbface/landmark_crop")

        v_ladmk_proj = sorted(jpath_ladmk_proj.glob("*.txt"))
        v_ladmk_crop = sorted(jpath_ladmk_crop.glob("*.txt"))
        print(f"contain {len(v_ladmk_proj)} / {len(v_ladmk_crop)} landmark proj/crop in {str(path_data)}")

        if len(v_ladmk_proj) != len(v_frames) or len(v_ladmk_crop) != len(v_frames):
            print(f"\t ERROR && Continue: This video missing landmark, {str(jpath_ladmk_proj)} ")
            # continue
        else:
            
            ladmk_proj_all = np.zeros((len(v_frames), 68, 2), np.float32)
            ladmk_crop_all = np.zeros((len(v_frames), 68, 2), np.float32)
            
            for k, iframe in enumerate(v_frames):
                fid = int(iframe.stem[-4:])
                    
                # load data
                kpath_proj = jpath_ladmk_proj.joinpath(f"frame-{fid:04d}.txt")
                kpath_crop = jpath_ladmk_crop.joinpath(f"frame-{fid:04d}.txt")

                ladmk_proj_all[k, :, :] = np.loadtxt(str(kpath_proj))   # (68, 2)
                ladmk_crop_all[k, :, :] = np.loadtxt(str(kpath_crop))   # (68, 2)

            # save combine data
            jpath_proj_save = path_data.joinpath("video/landmark_rgbface_proj.npy")   
            jpath_crop_save = path_data.joinpath("video/landmark_rgbface_crop.npy")   
            jpath_crop_save.parent.mkdir(parents=True, exist_ok=True)
            np.save(str(jpath_proj_save), ladmk_proj_all)
            np.save(str(jpath_crop_save), ladmk_crop_all)
            
            print(f"combine the coeff: {str(jpath_ladmk_proj)} --> {str(jpath_proj_save)}")
        # combine 3dmm of RGBFaceNet   
    
    if COMBINE_COEFF_RBGFACE:
        jpath_coeff = path_data.joinpath("rgbface/coeff")

        v_coeff = sorted(jpath_coeff.glob("*.txt"))
        print(f"contain {len(v_coeff)} reconstruction/tans in {str(path_data)}")

        if len(v_coeff) != len(v_frames):
            print(f"\t ERROR && Continue: This video missing coeff, {str(jpath_coeff)} ")
            # continue
        else:
            coeff_all = np.zeros((len(v_frames), 61), np.float32)  # 55 + 3 + 3
            
            for k, iframe in enumerate(v_frames):
                fid = int(iframe.stem[-4:])
                    
                # load data
                kpath_coeff = jpath_coeff.joinpath(f"frame-{fid:04d}.txt")

                kk_coeff = np.loadtxt(str(kpath_coeff)).reshape(-1)

                kk_coeff_expr = kk_coeff[100:155]     # (55, ) , identity-100, expression-55      ==> (55, )
                kk_coeff_rts  = kk_coeff[155:161]     # (6)    , rotation-3, translation-3, scale-1  ==> (6, )

                coeff_all[k, :] = np.r_[kk_coeff_expr, kk_coeff_rts]                  #  (61, )        expression-55 + rotation-3 + trans-3

            # save combine data
            jpath_coeff_save = path_data.joinpath("video/coeff_rgbface.npy")   
            jpath_coeff_save.parent.mkdir(parents=True, exist_ok=True)
            np.save(str(jpath_coeff_save), coeff_all)
            print(f"combine the coeff: {str(jpath_coeff)} --> {str(jpath_coeff_save)}")